{
 "cells": [
  {
   "cell_type": "markdown",
   "source": [
    "## MLP模型"
   ],
   "metadata": {
    "collapsed": false
   },
   "id": "64e859d6050fcc5f"
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "id": "initial_id",
   "metadata": {
    "collapsed": true,
    "ExecuteTime": {
     "end_time": "2025-05-15T07:04:40.064456Z",
     "start_time": "2025-05-15T07:04:33.123831Z"
    }
   },
   "outputs": [],
   "source": [
    "import torch\n",
    "import torch.nn as nn\n",
    "import torch.optim as optim\n",
    "import torchvision.transforms as transforms\n",
    "import os\n",
    "import numpy as np\n",
    "import time\n",
    "from sklearn.metrics import accuracy_score, precision_recall_fscore_support, confusion_matrix, classification_report\n",
    "import matplotlib.pyplot as plt\n",
    "import seaborn as sns\n",
    "import pickle"
   ]
  },
  {
   "cell_type": "code",
   "outputs": [],
   "source": [
    "# 数据集路径\n",
    "data_dir = r'D:\\Machine_learning\\jiqixuexi\\cifar-10-python\\cifar-10-batches-py'\n",
    "\n",
    "# 定义device变量，优先使用GPU\n",
    "device = torch.device(\"cuda\" if torch.cuda.is_available() else \"cpu\")\n",
    "\n",
    "# 自定义数据集加载函数\n",
    "class CIFAR10Dataset(torch.utils.data.Dataset):\n",
    "    def __init__(self, data, labels, transform=None):\n",
    "        self.data = data\n",
    "        self.labels = labels\n",
    "        self.transform = transform\n",
    "\n",
    "    def __len__(self):\n",
    "        return len(self.data)\n",
    "\n",
    "    def __getitem__(self, idx):\n",
    "        img = self.data[idx]\n",
    "        label = self.labels[idx]\n",
    "\n",
    "        if self.transform:\n",
    "            img = self.transform(img)\n",
    "\n",
    "        return img, label\n",
    "\n",
    "# 加载CIFAR-10数据集\n",
    "def load_cifar10_data(data_dir):\n",
    "    # 加载训练数据\n",
    "    train_data = []\n",
    "    train_labels = []\n",
    "    for i in range(1, 6):\n",
    "        batch_path = os.path.join(data_dir, f'data_batch_{i}')\n",
    "        with open(batch_path, 'rb') as file:\n",
    "            batch = pickle.load(file, encoding='latin1')\n",
    "            train_data.append(batch['data'])\n",
    "            train_labels.extend(batch['labels'])\n",
    "    train_data = np.array(train_data).reshape(50000, 3, 32, 32).transpose(0, 2, 3, 1)\n",
    "    train_labels = np.array(train_labels)\n",
    "\n",
    "    # 加载测试数据\n",
    "    test_batch_path = os.path.join(data_dir, 'test_batch')\n",
    "    with open(test_batch_path, 'rb') as file:\n",
    "        test_batch = pickle.load(file, encoding='latin1')\n",
    "    test_data = test_batch['data'].reshape(10000, 3, 32, 32).transpose(0, 2, 3, 1)\n",
    "    test_labels = np.array(test_batch['labels'])\n",
    "\n",
    "    return (train_data, train_labels), (test_data, test_labels)\n",
    "\n",
    "# 调用函数加载数据\n",
    "(train_data, train_labels), (test_data, test_labels) = load_cifar10_data(data_dir)\n",
    "\n",
    "# 数据预处理与增强\n",
    "transform_train = transforms.Compose([\n",
    "    transforms.ToTensor(),\n",
    "    transforms.Normalize((0.5, 0.5, 0.5), (0.5, 0.5, 0.5)),\n",
    "])\n",
    "\n",
    "transform_test = transforms.Compose([\n",
    "    transforms.ToTensor(),\n",
    "    transforms.Normalize((0.5, 0.5, 0.5), (0.5, 0.5, 0.5)),\n",
    "])\n",
    "\n",
    "# 自定义数据集\n",
    "train_dataset = CIFAR10Dataset(train_data, train_labels, transform=transform_train)\n",
    "test_dataset = CIFAR10Dataset(test_data, test_labels, transform=transform_test)\n",
    "\n",
    "# 数据加载器\n",
    "train_loader = torch.utils.data.DataLoader(train_dataset, batch_size=64, shuffle=True, num_workers=2)\n",
    "test_loader = torch.utils.data.DataLoader(test_dataset, batch_size=64, shuffle=False, num_workers=2)"
   ],
   "metadata": {
    "collapsed": false,
    "ExecuteTime": {
     "end_time": "2025-05-15T07:15:19.850778Z",
     "start_time": "2025-05-15T07:15:19.457093Z"
    }
   },
   "id": "7358d1171615d58",
   "execution_count": 2
  },
  {
   "cell_type": "code",
   "outputs": [],
   "source": [
    "# 定义MLP模型\n",
    "class MLP(nn.Module):\n",
    "    def __init__(self):\n",
    "        super(MLP, self).__init__()\n",
    "        self.fc1 = nn.Linear(32 * 32 * 3, 512)\n",
    "        self.fc2 = nn.Linear(512, 256)\n",
    "        self.fc3 = nn.Linear(256, 10)\n",
    "        self.dropout = nn.Dropout(0.5)\n",
    "\n",
    "    def forward(self, x):\n",
    "        x = x.view(x.size(0), -1)  # 将图像展平为一维向量\n",
    "        x = torch.relu(self.fc1(x))\n",
    "        x = self.dropout(x)\n",
    "        x = torch.relu(self.fc2(x))\n",
    "        x = self.dropout(x)\n",
    "        x = self.fc3(x)\n",
    "        return x\n",
    "\n",
    "# 初始化模型、损失函数和优化器，并将模型移动到device\n",
    "model = MLP().to(device)\n",
    "criterion = nn.CrossEntropyLoss()\n",
    "optimizer = optim.Adam(model.parameters(), lr=0.001)\n",
    "\n",
    "# 训练模型\n",
    "start_time = time.time()\n",
    "num_epochs = 20\n",
    "for epoch in range(num_epochs):\n",
    "    model.train()\n",
    "    epoch_loss = 0.0\n",
    "    for i, (images, labels) in enumerate(train_loader):\n",
    "        images = images.to(device)\n",
    "        labels = labels.to(device)\n",
    "        optimizer.zero_grad()\n",
    "        outputs = model(images)\n",
    "        loss = criterion(outputs, labels)\n",
    "        epoch_loss += loss.item()\n",
    "        loss.backward()\n",
    "        optimizer.step()\n",
    "    print(f'Epoch [{epoch+1}/{num_epochs}], Loss: {epoch_loss / len(train_loader):.4f}')\n",
    "\n",
    "train_time = time.time() - start_time\n",
    "print(f\"训练时间: {train_time:.2f}秒\")\n",
    "# 测试模型\n",
    "model.eval()\n",
    "with torch.no_grad():\n",
    "    correct = 0\n",
    "    total = 0\n",
    "    all_labels = []\n",
    "    all_preds = []\n",
    "    for images, labels in test_loader:\n",
    "        images = images.to(device)\n",
    "        labels = labels.to(device)\n",
    "        outputs = model(images)\n",
    "        _, predicted = torch.max(outputs.data, 1)\n",
    "        total += labels.size(0)\n",
    "        # 将布尔 张量 转换为整数类型，然后计算总和\n",
    "        correct += (predicted == labels).type(torch.int).sum().item()\n",
    "        all_labels.extend(labels.cpu().numpy())\n",
    "        all_preds.extend(predicted.cpu().numpy())\n",
    "\n",
    "    test_acc = correct / total\n",
    "    print(f\"测试准确率: {test_acc:.4f}\")"
   ],
   "metadata": {
    "collapsed": false,
    "is_executing": true,
    "ExecuteTime": {
     "start_time": "2025-05-15T07:15:21.338887Z"
    }
   },
   "id": "2682641ce86949dc"
  },
  {
   "cell_type": "code",
   "outputs": [],
   "source": [
    "# 计算评估指标\n",
    "accuracy = accuracy_score(all_labels, all_preds)\n",
    "precision, recall, f1, _ = precision_recall_fscore_support(all_labels, all_preds, average='weighted')\n",
    "conf_matrix = confusion_matrix(all_labels, all_preds)\n",
    "class_report = classification_report(all_labels, all_preds)\n",
    "\n",
    "# 输出结果\n",
    "print(\"Accuracy:\", accuracy)\n",
    "print(\"Precision:\", precision)\n",
    "print(\"Recall:\", recall)\n",
    "print(\"F1-score:\", f1)\n",
    "print(\"Confusion Matrix:\\n\", conf_matrix)\n",
    "print(\"Classification Report:\\n\", class_report)"
   ],
   "metadata": {
    "collapsed": false
   },
   "id": "2890b89c224b4f0c"
  },
  {
   "cell_type": "code",
   "outputs": [],
   "source": [
    "# 绘制混淆矩阵热力图\n",
    "plt.figure(figsize=(12, 8))\n",
    "sns.heatmap(conf_matrix, annot=True, fmt='d', cmap='Purple')\n",
    "plt.title(\"MLP - Confusion Matrix (CIFAR-10)\")\n",
    "plt.xlabel(\"Predicted Label\")\n",
    "plt.ylabel(\"True Label\")\n",
    "plt.show()"
   ],
   "metadata": {
    "collapsed": false
   },
   "id": "fc3b39d0e2cac03e"
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 2
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython2",
   "version": "2.7.6"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}
